package cn.bugstack.xfg.dev.tech.test;


import com.alibaba.fastjson.JSON;
import jakarta.annotation.Resource;
import lombok.extern.java.Log;
import lombok.extern.slf4j.Slf4j;
import org.junit.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
import org.junit.runner.RunWith;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @Author caoyi
 * @description:
 * @Date 2025/7/23 16:16
 */
@SpringBootTest
@Slf4j
@RunWith(SpringRunner.class)
public class RAGTest {
    @Resource
    private OllamaChatClient ollamaChatClient;
    @Resource
    private TokenTextSplitter tokenTextSplitter;
    @Resource
    private SimpleVectorStore simpleVectorStore;
    @Resource
    private PgVectorStore pgVectorStore;

    @Test
    public void upload(){
        TikaDocumentReader reader = new TikaDocumentReader("./data/file.text");
        List<Document> documents = reader.get();
        List<Document> documentSplitterList = tokenTextSplitter.apply(documents);//按token划分成一块一块的
        documentSplitterList.forEach(doc -> doc.getMetadata().put("knowledge", "cy知识库"));

        pgVectorStore.accept(documentSplitterList);
        log.info("上传完成");

    }

    @Test
    public void chat(){
        String message = "曹艺，哪年出生";

        String SYSTEM_PROMPT = """
                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
                If unsure, simply state that you don't know.
                Another thing you need to note is that your reply must be in Chinese!
                DOCUMENTS:
                    {documents}
                """;

        //向量库中查找对象
        SearchRequest searchRequest = SearchRequest.query(message).withTopK(5).withFilterExpression("knowledge=='cy知识库'");

        //执行查找
        List<Document> documents = pgVectorStore.similaritySearch(searchRequest);
        String collect = documents.stream().map(Document::getContent).collect(Collectors.joining());//合成一个string


        Message ragmsg = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", collect));

        ArrayList<Message> messages = new ArrayList<>();
        messages.add(new UserMessage(message));
        messages.add(ragmsg);

        ChatResponse call = ollamaChatClient.call(new Prompt(messages, OllamaOptions.create().withModel("deepseek-r1:1.5b")));

        log.info("测试结果:{}", JSON.toJSONString(call));
    }


}
